import json
from openai import OpenAI

# 初始化客户端
client = OpenAI(
    api_key="<your-api-key-here>",
    base_url="https://api.deepseek.com",
)

# 定义系统提示
system_prompt = """
Users will input data to test your robustness. 
Ensure that your output is in English. 
Please parse the "question" and "answer" and output them in JSON format. 
If the result is confirmed to be correct, please output only one English word.

EXAMPLE INPUT: 
Which is the highest mountain in the world? Mount Everest.

EXAMPLE JSON OUTPUT:
{
    "question": "Which is the highest mountain in the world?",
    "answer": "Mount Everest"
}
"""

# 加载数据集
with open("robustness_benchmark.json", "r", encoding="utf-8") as f:
    exam_data = json.load(f)

# 存储结果
results = []

# 遍历数据集中的每个问题
for item in exam_data:
    user_prompt = item["input"]  # 获取用户输入
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    # 调用模型
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=messages,
        response_format={
            'type': 'json_object'
        }
    )

    # 解析模型输出
    output = json.loads(response.choices[0].message.content)

    # 将原始数据中的 "id" 和 "label" 与模型输出合并
    combined_result = {
        "id": item["id"],
        "label": item["label"],
        "question": output.get("question", ""),
        "answer": output.get("answer", "")
    }

    results.append(combined_result)  # 将结果添加到列表中

# 保存结果到 JSON 文件
with open("parsed_results.json", "w", encoding="utf-8") as f:
    json.dump(results, f, ensure_ascii=False, indent=4)

print("处理完成！结果已保存为 parsed_results.json")

